#!/usr/bin/env python
# coding=utf-8



# change: update receive traffic kernel interface
import argparse
from collections import defaultdict, namedtuple
import functools
from subprocess import call
import sys
from time import sleep
from bcc import BPF
from socket import inet_ntop, AF_INET
from struct import pack
from queue import Queue

BPF_TXT = '''
#include <uapi/linux/ptrace.h>
#include <net/sock.h>
struct data_t{
    u32 pid;
    u32 addr;
    u16 sport;
    u16 dport;
};
BPF_HASH(tcp_send_bytes, struct data_t);
BPF_HASH(tcp_recv_bytes, struct data_t);
BPF_HASH(udp_send_bytes, struct data_t);
BPF_HASH(udp_recv_bytes, struct data_t);
BPF_HASH(tcp_send_store, u32,struct sock *);
BPF_HASH(udp_send_store, u32,struct sock *);
BPF_HASH(udp_recv_store, u32,struct sock *);



int tcp_sendmsgEnt(struct pt_regs *ctx, struct sock *sk,
    struct msghdr *msg, size_t size)
{
    u32 tid = bpf_get_current_pid_tgid();
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    ARG_PID
    tcp_send_store.update(&tid,&sk);
    return 0;
}

static int tcp_sendGet(int size){
    u32 tid = bpf_get_current_pid_tgid();
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    ARG_PID
    struct sock ** sockpp;
    sockpp = tcp_send_store.lookup(&tid);
    if (sockpp==0){
        return 0;
    }
    struct sock * sk = *sockpp;
    u16 dport = 0, family = sk->__sk_common.skc_family;

    if (family == AF_INET) {
        struct data_t ipv4_key = {.pid=pid};
        ipv4_key.addr = sk->__sk_common.skc_daddr;
        ipv4_key.sport = sk->__sk_common.skc_num;
        dport = sk->__sk_common.skc_dport;
        ipv4_key.dport = ntohs(dport);
        tcp_send_bytes.increment(ipv4_key, size);
    }
    tcp_send_store.delete(&tid);
    return 0;
}


int tcp_sendmsgRet(struct pt_regs *ctx){

    int size = PT_REGS_RC(ctx);
    if (size>0){
        return tcp_sendGet(size);
    }
    return 0;

}

int tcp_cleanup_rbufEnt(struct pt_regs *ctx, struct sock *sk, int copied)
{
    if (copied<0) return 0; // 提前处理负数
    u32 tid = bpf_get_current_pid_tgid();
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    ARG_PID
    u16 dport = 0, family = sk->__sk_common.skc_family;
    if (family == AF_INET) {
        struct data_t ipv4_key = {.pid=pid};
        ipv4_key.addr = sk->__sk_common.skc_rcv_saddr;
        ipv4_key.sport = sk->__sk_common.skc_num;
        dport = sk->__sk_common.skc_dport;
        ipv4_key.dport = ntohs(dport);
        tcp_recv_bytes.increment(ipv4_key, copied);
    }
    return 0;
}



int udp_sendmsgEnt(struct pt_regs *ctx, struct sock *sk,
    struct msghdr *msg, size_t size)
{
    u32 tid = bpf_get_current_pid_tgid();
    u32 pid = bpf_get_current_pid_tgid()>> 32;
    ARG_PID
    udp_send_store.update(&tid,&sk);
    return 0;
}

static int udp_sendGet(int size){
    u32 tid = bpf_get_current_pid_tgid();
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    ARG_PID
    struct sock ** sockpp ;
    sockpp= udp_send_store.lookup(&tid);
    if (sockpp==0){
        return 0;
    }
    struct sock * sk = *sockpp;
    u16 dport = 0, family = sk->__sk_common.skc_family;

    if (family == AF_INET) {
        struct data_t ipv4_key = {.pid=pid};
        ipv4_key.addr = sk->__sk_common.skc_daddr;
        ipv4_key.sport = sk->__sk_common.skc_num;
        dport = sk->__sk_common.skc_dport;
        ipv4_key.dport = ntohs(dport);
        udp_send_bytes.increment(ipv4_key, size);
    }
    udp_send_store.delete(&tid);
    return 0;
}

int udp_sendmsgRet(struct pt_regs *ctx){

    int size = PT_REGS_RC(ctx);
    if (size>0){
        return udp_sendGet(size);
    }
    return 0;

}


int udp_recvmsgEnt(struct pt_regs *ctx, struct sock *sk,  struct msghdr *msg, size_t len,int nonblock,int flags, int *addr_len)
{
    u32 tid = bpf_get_current_pid_tgid();
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    ARG_PID
    udp_recv_store.update(&tid,&sk);
    return 0;
}


static int udp_recvGet(int size){
    u32 tid = bpf_get_current_pid_tgid();
    u32 pid = bpf_get_current_pid_tgid() >> 32;
    ARG_PID
    struct sock ** sockpp ;
    sockpp= udp_recv_store.lookup(&tid);
    if (sockpp==0){
        return 0;
    }
    struct sock * sk = *sockpp;
    u16 dport = 0, family = sk->__sk_common.skc_family;
    u64 *val, zero =0;

    if (family == AF_INET) {
        struct data_t ipv4_key = {.pid=pid};
        ipv4_key.addr = sk->__sk_common.skc_rcv_saddr;
        ipv4_key.sport = sk->__sk_common.skc_num;
        dport = sk->__sk_common.skc_dport;
        ipv4_key.dport = ntohs(dport);
        udp_recv_bytes.increment(ipv4_key, size);
    }
    udp_recv_store.delete(&tid);
    return 0;
}


int udp_recvmsgRet(struct pt_regs *ctx){

    int size = PT_REGS_RC(ctx);
    if (size>0){
        return udp_recvGet(size);
    }
    return 0;
}



'''


SessionKey = namedtuple('Session',[ 'pid','comm','sport', 'addr', 'dport',"type"])
valPair = namedtuple('record',[ 'pid','comm','sport', 'addr', 'dport',"type","byteVal"])
@functools.cache
def pid_to_comm(pid):
    try:
        comm = open("/proc/%s/comm" % pid, "r").read().rstrip()
        return comm
    except IOError:
        return str(pid)

maxbit=sys.maxsize
x86 = True
if maxbit>2**32:
    x86 = False

def range_check(num)->bool:
    val = int(num)
    maxPid = 2**32 if x86 else 32768 
    if val<1 or val > maxPid:
        return False
    return True
def get_ipv4_session_key(k,isTcp):
    return SessionKey(pid=k.pid,comm=pid_to_comm(k.pid),sport=k.sport, addr=inet_ntop(AF_INET, pack("I", k.addr)), dport=k.dport,type="tcp" if isTcp else "udp")
def get_netflow(k,byteVal,isTcp):
    return valPair(pid=k.pid,comm=pid_to_comm(k.pid),sport=k.sport, addr=inet_ntop(AF_INET, pack("I", k.addr)), dport=k.dport,type="tcp" if isTcp else "udp",byteVal=byteVal)

parser = argparse.ArgumentParser(
    description = "Summarize send and recv flow by host",
    formatter_class = argparse.RawDescriptionHelpFormatter,
)
parser.add_argument("-p", "--pid",default="", 
    help = "Trace this pid only")
parser.add_argument("-i","--interval", default=1, type=int,
    help = "output interval, in second (default 1)")
arg = parser.parse_args()

def clear():
    call("clear")

def parser():
    global BPF_TXT
    if arg.pid !="":
        if range_check(arg.pid):
            Fliter_PID = f"if( pid != {arg.pid} ){{ return 0; }}"
    else:
            Fliter_PID = ""
    BPF_TXT = BPF_TXT.replace("ARG_PID",Fliter_PID)
    if arg.interval!=1:
        BPF_TXT = BPF_TXT.replace("ARG_INTERVAL",str(arg.interval))


sumTuple = namedtuple("summary",["total","totalSend","totalRecv"])

class Result:
    def __init__(self) -> None:
        self.recvs =[]
        self.sends =[]
        self.totals = [0,0,0]
    def setRecv(self,v):
        self.recvs.append(v)
        self.totals[0] +=v.byteVal

    def setSend(self,v):
        self.sends.append(v)
        self.totals[1] +=v.byteVal
        
    def summary(self):
        self.totals[2] = self.totals[0] + self.totals[1]




def exporter(que:Queue,interval:int):
    global BPF_TXT
    Fliter_PID = ""
    BPF_TXT = BPF_TXT.replace("ARG_PID",Fliter_PID)
    b = BPF(text=BPF_TXT)
    htab_batch_ops = True if BPF.kernel_struct_has_field(b'bpf_map_ops',
        b'map_lookup_and_delete_batch') == 1 else False
    b.attach_kprobe(event="tcp_sendmsg",fn_name="tcp_sendmsgEnt")
    b.attach_kretprobe(event="tcp_sendmsg",fn_name="tcp_sendmsgRet")
    b.attach_kprobe(event="tcp_cleanup_rbuf",fn_name="tcp_cleanup_rbufEnt")
    b.attach_kprobe(event="udp_sendmsg",fn_name="udp_sendmsgEnt")
    b.attach_kretprobe(event="udp_sendmsg",fn_name="udp_sendmsgRet")
    b.attach_kprobe(event="udp_recvmsg",fn_name="udp_recvmsgEnt")
    b.attach_kretprobe(event="udp_recvmsg",fn_name="udp_recvmsgRet")

    tcp_send_bytes = b["tcp_send_bytes"]
    tcp_recv_bytes = b["tcp_recv_bytes"]
    udp_send_bytes = b["udp_send_bytes"]
    udp_recv_bytes = b["udp_recv_bytes"]
    exiting = False
    while  not exiting:
        try:
            sleep(interval)
        except KeyboardInterrupt:
            exiting = True
        record = Result()
        for k, v in tcp_send_bytes.items():
            # print(k.addr,k.dport,k.sport)
            record.setSend(get_netflow(k,v.value,True))
        if not htab_batch_ops:
            tcp_send_bytes.clear()

        for k,v in tcp_recv_bytes.items():
            record.setRecv(get_netflow(k,v.value,True))
        if not htab_batch_ops:
            tcp_recv_bytes.clear()

        for k, v in udp_send_bytes.items():
            # print(k.addr,k.dport,k.sport)
            record.setSend(get_netflow(k,v.value,False))
        if not htab_batch_ops:
            udp_send_bytes.clear()

        for k,v in udp_recv_bytes.items():
            record.setRecv(get_netflow(k,v.value,False))
        if not htab_batch_ops:
            udp_recv_bytes.clear()
        #lock.acquire()
        record.summary()
        que.put(record)
                

        


if __name__ == "__main__":
    parser()
    b = BPF(text=BPF_TXT)
    htab_batch_ops = True if BPF.kernel_struct_has_field(b'bpf_map_ops',
        b'map_lookup_and_delete_batch') == 1 else False
    b.attach_kprobe(event="tcp_sendmsg",fn_name="tcp_sendmsgEnt")
    b.attach_kretprobe(event="tcp_sendmsg",fn_name="tcp_sendmsgRet")
    b.attach_kprobe(event="tcp_cleanup_rbuf",fn_name="tcp_cleanup_rbufEnt")
    b.attach_kprobe(event="udp_sendmsg",fn_name="udp_sendmsgEnt")
    b.attach_kretprobe(event="udp_sendmsg",fn_name="udp_sendmsgRet")
    b.attach_kprobe(event="udp_recvmsg",fn_name="udp_recvmsgEnt")
    b.attach_kretprobe(event="udp_recvmsg",fn_name="udp_recvmsgRet")

    tcp_send_bytes = b["tcp_send_bytes"]
    tcp_recv_bytes = b["tcp_recv_bytes"]
    udp_send_bytes = b["udp_send_bytes"]
    udp_recv_bytes = b["udp_recv_bytes"]

    # header

    # output
    sumrecv = 0
    sumsend = 0
    sum_kb = 0
    i = 0
    exiting = False
    while i != 10 and not exiting:
        try:
            sleep(arg.interval)
        except KeyboardInterrupt:
            exiting = True
        clear()    
        print("%-10s %-12s %-10s %-10s %-10s %-10s %-10s %-21s %-21s %3s" % ("PID","COMM","RX_KB", "TX_KB", "RXSUM_KB", "TXSUM_KB", "SUM_KB", "LADDR", "RADDR","TYPE"))

        ipv4_throughput = defaultdict(lambda:[0,0])
        for k, v in tcp_send_bytes.items():
            # print(k.addr,k.dport,k.sport)
            key=get_ipv4_session_key(k,True)
            ipv4_throughput[key][0] = v.value
        if not htab_batch_ops:
            tcp_send_bytes.clear()

        for k,v in tcp_recv_bytes.items():
            key = get_ipv4_session_key(k,True)
            ipv4_throughput[key][1] = v.value
        if not htab_batch_ops:
            tcp_recv_bytes.clear()
        for k, v in udp_send_bytes.items():
            # print(k.addr,k.dport,k.sport)
            key=get_ipv4_session_key(k,False)
            ipv4_throughput[key][0] = v.value
        if not htab_batch_ops:
            udp_send_bytes.clear()

        for k,v in udp_recv_bytes.items():
            key = get_ipv4_session_key(k,False)
            ipv4_throughput[key][1] = v.value
        if not htab_batch_ops:
            udp_recv_bytes.clear()
        #lock.acquire()
        res =""
        if ipv4_throughput:
            for k, (send_bytes, recv_bytes) in sorted(ipv4_throughput.items(),
                key=lambda kv: sum(kv[1]),
                reverse=True):
                recv_bytes = int(recv_bytes / 1024)
                send_bytes = int(send_bytes / 1024)
                if recv_bytes==0 and send_bytes==0:
                    continue
                sumrecv += recv_bytes
                sumsend += send_bytes
                sum_kb = sumrecv + sumsend
                res +="%-10d %-12.12s %-10d %-10d %-10d %-10d %-10d %-21s %-21s %3s\n" % (k.pid,k.comm,recv_bytes, send_bytes, sumrecv, sumsend, sum_kb, 
                    "localhost:" + str(k.sport), 
                    k.addr + ":" + str(k.dport),k.type)

        print(f"\r{res}",flush=True)
        #lock.release()
        i += 1